from __future__ import print_function
from os.path import join
import argparse
import torch
import torch.nn as nn
import time
import math
import os
from torch.autograd import Variable
import numpy as np
from PIL import Image
from scipy.signal import convolve2d
from tqdm import tqdm
from torchvision.transforms import ToTensor
import numpy as np


def psnr(img1, img2):
    mse = torch.mean((img1 - img2) ** 2)
    if mse < 1.0e-10:
        return 100
    PIXEL_MAX = 1
    return 20 * math.log10(PIXEL_MAX / math.sqrt(mse))


def matlab_style_gauss2D(shape=(3, 3), sigma=0.5):
    """
    2D gaussian mask - should give the same result as MATLAB's
    fspecial('gaussian',[shape],[sigma])
    """
    m, n = [(ss - 1.) / 2. for ss in shape]
    y, x = np.ogrid[-m:m + 1, -n:n + 1]
    h = np.exp(-(x * x + y * y) / (2. * sigma * sigma))
    h[h < np.finfo(h.dtype).eps * h.max()] = 0
    sumh = h.sum()
    if sumh != 0:
        h /= sumh
    return h


def filter2(x, kernel, mode='same'):
    return convolve2d(x, np.rot90(kernel, 2), mode=mode)


def compute_ssim(im1, im2, k1=0.01, k2=0.03, win_size=11, L=255):
    if not im1.shape == im2.shape:
        raise ValueError("Input Imagees must have the same dimensions")
    if len(im1.shape) > 2:
        raise ValueError("Please input the images with 1 channel")

    M, N = im1.shape
    C1 = (k1 * L) ** 2
    C2 = (k2 * L) ** 2
    window = matlab_style_gauss2D(shape=(win_size, win_size), sigma=1.5)
    window = window / np.sum(np.sum(window))

    if im1.dtype == np.uint8:
        im1 = np.double(im1)
    if im2.dtype == np.uint8:
        im2 = np.double(im2)

    mu1 = filter2(im1, window, 'valid')
    mu2 = filter2(im2, window, 'valid')
    mu1_sq = mu1 * mu1
    mu2_sq = mu2 * mu2
    mu1_mu2 = mu1 * mu2
    sigma1_sq = filter2(im1 * im1, window, 'valid') - mu1_sq
    sigma2_sq = filter2(im2 * im2, window, 'valid') - mu2_sq
    sigmal2 = filter2(im1 * im2, window, 'valid') - mu1_mu2

    ssim_map = ((2 * mu1_mu2 + C1) * (2 * sigmal2 + C2)) / \
        ((mu1_sq + mu2_sq + C1) * (sigma1_sq + sigma2_sq + C2))

    return np.mean(np.mean(ssim_map))


total_psnr = 0
index = 1
total_ssim = 0
TESTSET = "Cn20"
HR_IMAGE_PATH = "./Data/HR-BGAN-"+TESTSET
SR_IMAGE_PATH = "./Data/"+TESTSET
images = os.listdir(HR_IMAGE_PATH)
images_loader = tqdm(images, desc="[calculating...]")
for image_name in images_loader:
    hr_path = os.path.join(HR_IMAGE_PATH, image_name)
    sr_path = os.path.join(SR_IMAGE_PATH, image_name)
    sr = Image.open(sr_path)
    hr = Image.open(hr_path)
    total_ssim += compute_ssim(np.array(hr.convert("L")),
                               np.array(sr.convert("L")))
    sr = Variable(ToTensor()(sr)).unsqueeze(0)
    hr = Variable(ToTensor()(hr)).unsqueeze(0)

    total_psnr += psnr(sr, hr)
    index += 1

print("avg_psnr:%f" % (total_psnr/(index-1)))
print("avg_ssim:%f" % (total_ssim/(index-1)))
